import sys
import sys
import requests
import tqdm
import click
import hashlib
import time
import os
import concurrent.futures
from threading import Lock, Thread
import re
from urllib.parse import urlparse, unquote

# 尝试导入SMB和WebDAV支持库
try:
    # 使用pysmb替代smbprotocol
    from smb.SMBConnection import SMBConnection
    from smb import smb_structs
    smb_structs.SUPPORT_SMB2 = True
    SMB_SUPPORTED = True
except Exception as e:
    SMB_SUPPORTED = False

# 尝试导入PySide6用于GUI界面
try:
    from PySide6.QtWidgets import (
        QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
        QLabel, QLineEdit, QPushButton, QFileDialog, QProgressBar,
        QTextEdit, QComboBox, QCheckBox, QFrame, QGridLayout, QGroupBox
    )
    from PySide6.QtCore import Qt, QThread, Signal, Slot
    from PySide6.QtGui import QFont, QIcon, QColor, QPalette, QTextCursor
    GUI_SUPPORTED = True
except ImportError:
    GUI_SUPPORTED = False

@click.command()
@click.option("-u", "--url", required=False, help="download url (http, https, smb, webdav)")
@click.option("-c", "--chunk-size", default=8192, help="chunk size for downloading, default: 8192 bytes")
@click.option("-o", "--output", default=None, help="output filename, default: use filename from URL")
@click.option("-t", "--threads", default=4, help="number of download threads, default: 4")
@click.option("--timeout", default=30, help="request timeout in seconds, default: 30")
@click.option("--username", default=None, help="username for authentication (for smb/webdav)")
@click.option("--password", default=None, help="password for authentication (for smb/webdav)")
@click.option("--use-cli", is_flag=True, help="Use command line interface instead of GUI")
def download(url, chunk_size, output, threads, timeout, username, password, use_cli=False):
    # 检查是否使用命令行模式或没有GUI支持
    if use_cli or not GUI_SUPPORTED:
        # 命令行模式 - 检查必需参数
        if not url and not use_cli:
            click.secho("Error: URL is required when not using GUI mode", fg="red")
            return
        
        # 执行命令行下载
        run_cli_download(url, chunk_size, output, threads, timeout, username, password)
    else:
        # GUI模式
        run_gui_download(url, chunk_size, output, threads, timeout, username, password)

def run_cli_download(url, chunk_size, output, threads, timeout, username, password):
    start_time = time.time()
    
    # Parse URL to determine protocol
    parsed_url = urlparse(url)
    protocol = parsed_url.scheme.lower()
    
    # Extract filename from URL or use specified output
    if output:
        # 检查指定路径是否为目录
        if os.path.isdir(output):
            # 如果是目录，则从URL提取文件名并与目录路径组合
            if protocol in ['http', 'https', 'webdav', 'dav']:
                filename_from_url = url.split("/")[-1]
            elif protocol == 'smb':
                # For SMB, extract the filename from the path
                path_parts = parsed_url.path.split('/')
                filename_from_url = path_parts[-1] if path_parts[-1] else path_parts[-2]
            else:
                filename_from_url = url.split("/")[-1]
            
            # Remove any query string from filename
            filename_from_url = filename_from_url.split('?')[0]
            
            # 确保文件名有效
            if not filename_from_url or len(filename_from_url) < 3:
                # 生成基于时间戳的文件名
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                filename_from_url = f"download_{timestamp}.bin"
            
            # 添加默认扩展名（如果没有）
            if '.' not in filename_from_url:
                filename_from_url += ".bin"
            
            # 组合目录路径和文件名
            filename = os.path.join(output, filename_from_url)
            click.secho(f"检测到目录路径，将文件保存为: {os.path.basename(filename)}", fg="blue")
        else:
            filename = output
    else:
        if protocol in ['http', 'https', 'webdav', 'dav']:
            filename = url.split("/")[-1]
        elif protocol == 'smb':
            # For SMB, extract the filename from the path
            path_parts = parsed_url.path.split('/')
            filename = path_parts[-1] if path_parts[-1] else path_parts[-2]
        else:
            filename = url.split("/")[-1]
    
    # Remove any query string from filename
    filename = filename.split('?')[0]
    
    # 初始化速度统计变量
    speeds = []
    max_speed = 0
    min_speed = float('inf')
    
    try:
        # Handle different protocols
        if protocol in ['http', 'https']:
            # HTTP/HTTPS download logic (existing code)
            user_agent = {"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"}
            
            # Create session for all requests
            session = requests.Session()
            session.headers.update(user_agent)
            session.timeout = timeout
            
            # Add authentication if provided
            if username and password:
                session.auth = (username, password)
            
            # First request to get file info
            with session.head(url) as response:
                response.raise_for_status()
                total_size = int(response.headers.get("content-length", 0))
                
                # Check if server supports range requests
                accept_ranges = response.headers.get("accept-ranges", "none")
                supports_ranges = accept_ranges.lower() == "bytes"
            
            # 定义进度回调函数以跟踪速度
            def cli_progress_callback(bytes_downloaded, total_bytes, speed=None):
                nonlocal speeds, max_speed, min_speed
                if speed:
                    speeds.append(speed)
                    max_speed = max(max_speed, speed)
                    min_speed = min(min_speed, speed)
            
            download_http_https(url, filename, chunk_size, threads, timeout, total_size, supports_ranges, session, cli_progress_callback)
            
        elif protocol == 'smb':
            # SMB download logic
            if not SMB_SUPPORTED:
                raise ImportError("SMB support not available. Install pysmb library with 'pip install pysmb'.")
            
            click.secho(f"Downloading from SMB: {url}", fg="blue")
            download_smb(url, filename, chunk_size, username, password, timeout, None)
            
        elif protocol in ['webdav', 'dav']:
            # WebDAV download logic
            click.secho(f"Downloading from WebDAV: {url}", fg="blue")
            # Convert webdav:// to http:// for requests library
            if protocol == 'webdav':
                http_url = url.replace('webdav://', 'http://')
            else:
                http_url = url.replace('dav://', 'http://')
            
            # Create session for WebDAV requests
            session = requests.Session()
            session.timeout = timeout
            
            # Add authentication if provided
            if username and password:
                session.auth = (username, password)
                
            # WebDAV uses HTTP methods, so we can use similar logic to HTTP
            try:
                # PROPFIND to get file info (WebDAV way)
                headers = {'Depth': '0'}
                with session.request('PROPFIND', http_url, headers=headers) as response:
                    response.raise_for_status()
                    # Extract content-length from response if possible
                    # For simplicity, we'll use a GET request to get content length
                    head_response = session.head(http_url)
                    total_size = int(head_response.headers.get("content-length", 0))
                    
                    # Check if server supports range requests
                    accept_ranges = head_response.headers.get("accept-ranges", "none")
                    supports_ranges = accept_ranges.lower() == "bytes"
                    
            except:
                # Fallback to standard HEAD request
                with session.head(http_url) as response:
                    response.raise_for_status()
                    total_size = int(response.headers.get("content-length", 0))
                    accept_ranges = response.headers.get("accept-ranges", "none")
                    supports_ranges = accept_ranges.lower() == "bytes"
            
            # Use HTTP download logic for WebDAV
            download_http_https(http_url, filename, chunk_size, threads, timeout, total_size, supports_ranges, session, progress_callback=None)
            
        else:
            raise ValueError(f"Unsupported protocol: {protocol}")
            
        end_time = time.time()
        click.secho(f"Download complete, time cost: {end_time - start_time:.2f} seconds", fg="green")
        
        # Calculate and display MD5 hash
        md5_hash = hashlib.md5()
        with open(filename, "rb") as f:
            while True:
                chunk = f.read(65536)
                if not chunk:
                    break
                md5_hash.update(chunk)
        click.secho(f"MD5 hash: {md5_hash.hexdigest()}", fg="green")
        
        # 显示速度统计
        if speeds:
            avg_speed = sum(speeds) / len(speeds)
            
            # 格式化速度显示
            if avg_speed < 1:
                avg_speed_str = f"{avg_speed * 1024:.2f} KB/s"
            else:
                avg_speed_str = f"{avg_speed:.2f} MB/s"
            
            if max_speed < 1:
                max_speed_str = f"{max_speed * 1024:.2f} KB/s"
            else:
                max_speed_str = f"{max_speed:.2f} MB/s"
            
            # 处理最小速度可能为无穷大的情况
            if min_speed == float('inf'):
                min_speed_str = "0.00 MB/s"
            elif min_speed < 1:
                min_speed_str = f"{min_speed * 1024:.2f} KB/s"
            else:
                min_speed_str = f"{min_speed:.2f} MB/s"
            
            click.secho(f"下载速度统计:", fg="blue")
            click.secho(f"  平均速度: {avg_speed_str}", fg="green")
            click.secho(f"  最高速度: {max_speed_str}", fg="green")
            click.secho(f"  最低速度: {min_speed_str}", fg="green")
        
    except Exception as e:
        click.secho(f"Error during download: {e}", fg="red")
        # Clean up partially downloaded file
        if os.path.exists(filename):
            os.remove(filename)
        return

# GUI相关类
class DownloadWorker(QThread):
    """下载工作线程类"""
    progress_updated = Signal(int, str, str)  # progress_value, total_size, speed
    download_complete = Signal(bool, str, str, str, str, str)  # success, error_msg, md5_hash, avg_speed, max_speed, min_speed
    log_message = Signal(str, str)  # message, level (info/error/warning)
    
    def __init__(self, url, output_file, chunk_size, threads, timeout, username, password):
        super().__init__()
        self.url = url
        self.output_file = output_file
        self.chunk_size = chunk_size
        self.threads = threads
        self.timeout = timeout
        self.username = username
        self.password = password
        self._stop_flag = False
        # 添加速度统计变量
        self.speeds = []
        self.max_speed = 0
        self.min_speed = float('inf')
    
    def stop(self):
        self._stop_flag = True
    
    def run(self):
        try:
            # Parse URL to determine protocol
            parsed_url = urlparse(self.url)
            protocol = parsed_url.scheme.lower()
            
            # 进度回调函数
            def progress_callback(bytes_downloaded, total_bytes, speed=None):
                if self._stop_flag:
                    raise InterruptedError("Download stopped by user")
                
                if total_bytes > 0:
                    progress_value = int((bytes_downloaded / total_bytes) * 100)
                else:
                    progress_value = 0
                
                # 将大数值转换为字符串，避免整数溢出问题
                total_bytes_str = str(total_bytes)
                
                # 记录速度数据
                if speed:
                    self.speeds.append(speed)
                    self.max_speed = max(self.max_speed, speed)
                    self.min_speed = min(self.min_speed, speed)
                
                speed_str = f"{speed:.2f} MB/s" if speed else "Calculating..."
                self.progress_updated.emit(progress_value, total_bytes_str, speed_str)
            
            # 根据协议调用相应的下载函数
            if protocol in ['http', 'https']:
                self.log_message.emit(f"开始下载 HTTP/HTTPS 文件: {self.url}", "info")
                # 创建会话
                user_agent = {"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"}
                session = requests.Session()
                session.headers.update(user_agent)
                session.timeout = self.timeout
                
                if self.username and self.password:
                    session.auth = (self.username, self.password)
                
                # 获取文件信息
                with session.head(self.url) as response:
                    response.raise_for_status()
                    total_size = int(response.headers.get("content-length", 0))
                    accept_ranges = response.headers.get("accept-ranges", "none")
                    supports_ranges = accept_ranges.lower() == "bytes"
                
                download_http_https(
                    self.url, self.output_file, self.chunk_size, 
                    self.threads, self.timeout, total_size, 
                    supports_ranges, session, progress_callback
                )
                
            elif protocol == 'smb':
                self.log_message.emit(f"开始下载 SMB 文件: {self.url}", "info")
                if not SMB_SUPPORTED:
                    raise ImportError("SMB support not available. Install pysmb library.")
                
                download_smb(
                    self.url, self.output_file, self.chunk_size, 
                    self.username, self.password, self.timeout, 
                    progress_callback
                )
                
            elif protocol in ['webdav', 'dav']:
                self.log_message.emit(f"开始下载 WebDAV 文件: {self.url}", "info")
                # 转换为HTTP URL
                http_url = self.url.replace('webdav://', 'http://').replace('dav://', 'http://')
                
                # 创建会话
                session = requests.Session()
                session.timeout = self.timeout
                
                if self.username and self.password:
                    session.auth = (self.username, self.password)
                
                # 获取文件信息
                try:
                    headers = {'Depth': '0'}
                    with session.request('PROPFIND', http_url, headers=headers) as response:
                        response.raise_for_status()
                        head_response = session.head(http_url)
                        total_size = int(head_response.headers.get("content-length", 0))
                        accept_ranges = head_response.headers.get("accept-ranges", "none")
                        supports_ranges = accept_ranges.lower() == "bytes"
                except:
                    with session.head(http_url) as response:
                        response.raise_for_status()
                        total_size = int(response.headers.get("content-length", 0))
                        accept_ranges = response.headers.get("accept-ranges", "none")
                        supports_ranges = accept_ranges.lower() == "bytes"
                
                download_http_https(
                    http_url, self.output_file, self.chunk_size, 
                    self.threads, self.timeout, total_size, 
                    supports_ranges, session, progress_callback
                )
                
            else:
                raise ValueError(f"不支持的协议: {protocol}")
            
            # 计算MD5
            try:
                md5_hash = hashlib.md5()
                with open(self.output_file, "rb") as f:
                    while True:
                        chunk = f.read(65536)
                        if not chunk:
                            break
                        md5_hash.update(chunk)
                md5_hash_value = md5_hash.hexdigest()
            except Exception as e:
                self.log_message.emit(f"计算MD5时出错: {str(e)}", "warning")
                md5_hash_value = "error"
            
            # 计算速度统计
            avg_speed_str = "0.00 MB/s"
            max_speed_str = "0.00 MB/s"
            min_speed_str = "0.00 MB/s"
            
            if self.speeds:
                avg_speed = sum(self.speeds) / len(self.speeds)
                # 格式化速度显示
                if avg_speed < 1:
                    avg_speed_str = f"{avg_speed * 1024:.2f} KB/s"
                else:
                    avg_speed_str = f"{avg_speed:.2f} MB/s"
                
                if self.max_speed < 1:
                    max_speed_str = f"{self.max_speed * 1024:.2f} KB/s"
                else:
                    max_speed_str = f"{self.max_speed:.2f} MB/s"
                
                # 处理最小速度可能为无穷大的情况
                if self.min_speed == float('inf'):
                    min_speed_str = "0.00 MB/s"
                elif self.min_speed < 1:
                    min_speed_str = f"{self.min_speed * 1024:.2f} KB/s"
                else:
                    min_speed_str = f"{self.min_speed:.2f} MB/s"
            
            self.log_message.emit(f"下载完成! MD5: {md5_hash_value}", "info")
            self.download_complete.emit(True, "下载成功", md5_hash_value, avg_speed_str, max_speed_str, min_speed_str)
            
        except OSError as e:
            # 特别处理操作系统错误，如权限问题
            error_msg = str(e)
            if "Permission denied" in error_msg:
                # 更友好的权限错误提示
                self.log_message.emit(f"下载错误: 权限被拒绝 - 无法写入到目标位置 '{self.output_file}'", "error")
                self.log_message.emit("提示: 请选择一个有写入权限的目录，如桌面或文档文件夹", "warning")
            else:
                self.log_message.emit(f"下载错误: 文件系统错误 - {error_msg}", "error")
            
            # 清理部分下载的文件
            if os.path.exists(self.output_file):
                try:
                    os.remove(self.output_file)
                except:
                    pass
            
            self.download_complete.emit(False, error_msg)
            
        except Exception as e:
            # 处理其他类型的错误
            self.log_message.emit(f"下载错误: {str(e)}", "error")
            
            # 清理部分下载的文件
            if os.path.exists(self.output_file):
                try:
                    os.remove(self.output_file)
                except:
                    pass
            
            self.download_complete.emit(False, str(e), "", "", "", "")

class DownloadManager(QMainWindow):
    """下载管理器GUI类"""
    def __init__(self, initial_url="", initial_output=None, initial_threads=4, initial_timeout=30):
        super().__init__()
        self.worker = None
        self.initial_url = initial_url
        self.initial_output = initial_output
        self.initial_threads = initial_threads
        self.initial_timeout = initial_timeout
        self.init_ui()
    
    def init_ui(self):
        # 设置窗口标题和大小
        self.setWindowTitle("高级下载管理器")
        self.resize(800, 600)
        
        # 设置字体
        font = QFont("微软雅黑", 10)
        self.setFont(font)
        
        # 创建主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(15, 15, 15, 15)
        main_layout.setSpacing(12)
        
        # 样式设置
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f0f0f0;
            }
            QGroupBox {
                background-color: #ffffff;
                border: 1px solid #d0d0d0;
                border-radius: 8px;
                margin-top: 10px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                background-color: #4a90e2;
                color: white;
                padding: 3px 15px;
                border-top-left-radius: 8px;
                border-top-right-radius: 8px;
            }
            QLineEdit, QComboBox {
                padding: 8px 10px;
                border: 1px solid #d0d0d0;
                border-radius: 4px;
                background-color: #ffffff;
            }
            QLineEdit:focus, QComboBox:focus {
                border-color: #4a90e2;
                background-color: #f8f9ff;
            }
            QPushButton {
                padding: 8px 15px;
                border: none;
                border-radius: 4px;
                background-color: #4a90e2;
                color: white;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #357abd;
            }
            QPushButton:pressed {
                background-color: #2968a3;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
            QProgressBar {
                border: 1px solid #d0d0d0;
                border-radius: 4px;
                background-color: #f0f0f0;
                text-align: center;
                height: 20px;
            }
            QProgressBar::chunk {
                background-color: #4a90e2;
                border-radius: 4px;
            }
            QTextEdit {
                border: 1px solid #d0d0d0;
                border-radius: 4px;
                background-color: #ffffff;
                font-family: Consolas, monospace;
                font-size: 9pt;
            }
        """)
        
        # URL输入区域
        url_group = QGroupBox("下载设置")
        url_layout = QGridLayout()
        url_group.setLayout(url_layout)
        
        url_layout.addWidget(QLabel("下载链接:"), 0, 0, Qt.AlignRight | Qt.AlignVCenter)
        self.url_input = QLineEdit(self.initial_url)
        self.url_input.setPlaceholderText("请输入下载链接 (http, https, smb, webdav)")
        self.url_input.setMinimumWidth(400)
        url_layout.addWidget(self.url_input, 0, 1, 1, 2)
        
        url_layout.addWidget(QLabel("保存路径:"), 1, 0, Qt.AlignRight | Qt.AlignVCenter)
        self.output_input = QLineEdit(self.initial_output or "")
        self.output_input.setPlaceholderText("请选择保存路径")
        url_layout.addWidget(self.output_input, 1, 1)
        
        self.browse_button = QPushButton("浏览...")
        self.browse_button.clicked.connect(self.browse_output)
        url_layout.addWidget(self.browse_button, 1, 2)
        
        # 高级设置区域
        advanced_group = QGroupBox("高级设置")
        advanced_layout = QGridLayout()
        advanced_group.setLayout(advanced_layout)
        
        advanced_layout.addWidget(QLabel("线程数:"), 0, 0, Qt.AlignRight | Qt.AlignVCenter)
        self.threads_combo = QComboBox()
        for i in [1, 2, 4, 8, 16, 32]:
            self.threads_combo.addItem(str(i))
        self.threads_combo.setCurrentText(str(self.initial_threads))
        advanced_layout.addWidget(self.threads_combo, 0, 1)
        
        advanced_layout.addWidget(QLabel("超时设置 (秒):"), 0, 2, Qt.AlignRight | Qt.AlignVCenter)
        self.timeout_input = QLineEdit(str(self.initial_timeout))
        self.timeout_input.setMaximumWidth(100)
        advanced_layout.addWidget(self.timeout_input, 0, 3)
        
        advanced_layout.addWidget(QLabel("用户名:"), 1, 0, Qt.AlignRight | Qt.AlignVCenter)
        self.username_input = QLineEdit()
        self.username_input.setPlaceholderText("SMB/WebDAV认证用户名")
        advanced_layout.addWidget(self.username_input, 1, 1)
        
        advanced_layout.addWidget(QLabel("密码:"), 1, 2, Qt.AlignRight | Qt.AlignVCenter)
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.Password)
        self.password_input.setPlaceholderText("SMB/WebDAV认证密码")
        advanced_layout.addWidget(self.password_input, 1, 3)
        
        # 下载控制区域
        control_layout = QHBoxLayout()
        
        self.start_button = QPushButton("开始下载")
        self.start_button.clicked.connect(self.start_download)
        self.start_button.setMinimumWidth(100)
        control_layout.addWidget(self.start_button)
        
        self.stop_button = QPushButton("停止下载")
        self.stop_button.clicked.connect(self.stop_download)
        self.stop_button.setMinimumWidth(100)
        self.stop_button.setEnabled(False)
        control_layout.addWidget(self.stop_button)
        
        control_layout.addStretch()
        
        # 进度条区域
        progress_group = QGroupBox("下载进度")
        progress_layout = QVBoxLayout()
        progress_group.setLayout(progress_layout)
        
        progress_info_layout = QHBoxLayout()
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        self.progress_bar.setTextVisible(True)
        progress_info_layout.addWidget(self.progress_bar)
        
        self.speed_label = QLabel("0 MB/s")
        progress_info_layout.addWidget(self.speed_label)
        progress_layout.addLayout(progress_info_layout)
        
        # 日志区域
        log_group = QGroupBox("下载日志")
        log_layout = QVBoxLayout()
        log_group.setLayout(log_layout)
        
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setMinimumHeight(200)
        log_layout.addWidget(self.log_text)
        
        # 添加到主布局
        main_layout.addWidget(url_group)
        main_layout.addWidget(advanced_group)
        main_layout.addLayout(control_layout)
        main_layout.addWidget(progress_group)
        main_layout.addWidget(log_group, 1)
    
    def browse_output(self):
        filename, _ = QFileDialog.getSaveFileName(
            self, "保存文件", "", "所有文件 (*)"
        )
        if filename:
            self.output_input.setText(filename)
    
    def start_download(self):
        self.log_text.clear()
        url = self.url_input.text().strip()
        if not url:
            self.add_log("请输入下载链接", "error")
            return
        
        output_file = self.output_input.text().strip()
        if not output_file:
            # 从URL提取文件名
            try:
                parsed_url = urlparse(url)
                output_file = url.split("/")[-1].split('?')[0]
                
                # 确保文件名有效
                if not output_file or len(output_file) < 3:
                    # 生成基于时间戳的文件名
                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                    output_file = f"download_{timestamp}.bin"
                
                # 添加默认扩展名（如果没有）
                if '.' not in output_file:
                    # 根据URL协议或内容类型添加合理的扩展名
                    output_file += ".bin"
            except:
                # 使用时间戳确保文件名唯一且有效
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                output_file = f"download_{timestamp}.bin"
        else:
            # 验证输出路径是否有效
            try:
                # 检查目录是否存在且可写
                dir_path = os.path.dirname(output_file)
                if dir_path and not os.path.exists(dir_path):
                    try:
                        os.makedirs(dir_path)
                    except Exception as e:
                        self.add_log(f"无法创建目录: {str(e)}", "error")
                        return
                
                # 检查指定路径是否为目录
                if os.path.isdir(output_file):
                    # 如果是目录，则从URL提取文件名并与目录路径组合
                    try:
                        parsed_url = urlparse(url)
                        # 从URL提取文件名
                        filename_from_url = url.split("/")[-1].split('?')[0]
                        
                        # 确保文件名有效
                        if not filename_from_url or len(filename_from_url) < 3:
                            # 生成基于时间戳的文件名
                            timestamp = time.strftime("%Y%m%d_%H%M%S")
                            filename_from_url = f"download_{timestamp}.bin"
                        
                        # 添加默认扩展名（如果没有）
                        if '.' not in filename_from_url:
                            # 根据URL协议或内容类型添加合理的扩展名
                            filename_from_url += ".bin"
                        
                        # 组合目录路径和文件名
                        output_file = os.path.join(output_file, filename_from_url)
                        self.add_log(f"检测到目录路径，将文件保存为: {os.path.basename(output_file)}", "info")
                    except Exception as e:
                        self.add_log(f"处理文件名时出错: {str(e)}", "error")
                        return
            except Exception as e:
                self.add_log(f"无效的保存路径: {str(e)}", "error")
                return
        
        try:
            threads = int(self.threads_combo.currentText())
            timeout = int(self.timeout_input.text())
        except ValueError:
            self.add_log("线程数和超时设置必须是数字", "error")
            return
        
        username = self.username_input.text()
        password = self.password_input.text()
        chunk_size = 8192  # 默认块大小
        
        # 禁用开始按钮，启用停止按钮
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        
        # 创建并启动下载线程
        self.worker = DownloadWorker(
            url, output_file, chunk_size, threads, timeout, username, password
        )
        self.worker.progress_updated.connect(self.update_progress)
        self.worker.download_complete.connect(self.download_finished)
        self.worker.log_message.connect(self.add_log)
        self.worker.start()
    
    def stop_download(self):
        if self.worker and self.worker.isRunning():
            self.worker.stop()
            self.add_log("正在停止下载...", "warning")
    
    def update_progress(self, progress, total_size_str, speed):
        self.progress_bar.setValue(progress)
        self.speed_label.setText(speed)
        
        # 如果需要显示总大小，可以添加一个大小标签
        # 例如: self.total_size_label.setText(f"总大小: {total_size_str} 字节")
    
    def download_finished(self, success, message, md5_hash="", avg_speed="", max_speed="", min_speed=""):
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        
        if success:
            self.add_log(f"下载完成: {message}", "info")
            # 显示速度统计信息
            if avg_speed:
                self.add_log(f"平均速度: {avg_speed}", "info")
            if max_speed:
                self.add_log(f"最高速度: {max_speed}", "info")
            if min_speed:
                self.add_log(f"最低速度: {min_speed}", "info")
        else:
            self.add_log(f"下载失败: {message}", "error")
    
    def add_log(self, message, level="info"):
        timestamp = time.strftime("%H:%M:%S")
        
        if level == "error":
            prefix = "[错误]"
            color = "#ff4444"
        elif level == "warning":
            prefix = "[警告]"
            color = "#ffaa00"
        else:
            prefix = "[信息]"
            color = "#000000"
        
        log_line = f"<font color='{color}'>{timestamp} {prefix} {message}</font><br>"
        
        self.log_text.insertHtml(log_line)
        self.log_text.moveCursor(QTextCursor.End)

def run_gui_download(url, chunk_size, output, threads, timeout, username, password):
    """运行GUI版本的下载器"""
    app = QApplication(sys.argv)
    app.setApplicationName("高级下载管理器")
    app.setStyle("Fusion")
    
    # 设置全局调色板
    palette = app.palette()
    palette.setColor(QPalette.Window, QColor(240, 240, 240))
    palette.setColor(QPalette.WindowText, QColor(0, 0, 0))
    palette.setColor(QPalette.Base, QColor(255, 255, 255))
    palette.setColor(QPalette.AlternateBase, QColor(240, 240, 240))
    palette.setColor(QPalette.Highlight, QColor(74, 144, 226))
    palette.setColor(QPalette.HighlightedText, QColor(255, 255, 255))
    app.setPalette(palette)
    
    # 创建下载管理器窗口
    download_manager = DownloadManager(url, output, threads, timeout)
    
    # 如果提供了用户名密码，设置它们
    if username:
        download_manager.username_input.setText(username)
    if password:
        download_manager.password_input.setText(password)
    
    # 显示窗口
    download_manager.show()
    
    # 运行应用
    sys.exit(app.exec())

def download_http_https(url, filename, chunk_size, threads, timeout, total_size, supports_ranges, session, progress_callback=None):
    """Download file using HTTP/HTTPS protocol"""
    # 初始化进度跟踪变量
    bytes_downloaded = 0
    start_time = time.time()
    last_update_time = start_time
    last_update_bytes = 0
    
    # 是否需要显示命令行进度条
    use_cli_progress = progress_callback is None
    
    # 进度更新函数
    def update_progress(chunk_size_downloaded):
        nonlocal bytes_downloaded, last_update_time, last_update_bytes
        bytes_downloaded += chunk_size_downloaded
        
        # 计算下载速度
        current_time = time.time()
        time_diff = current_time - last_update_time
        
        # 每0.5秒或下载了足够的数据时更新
        if time_diff > 0.5 or bytes_downloaded >= total_size:
            bytes_diff = bytes_downloaded - last_update_bytes
            speed = (bytes_diff / time_diff) / (1024 * 1024)  # MB/s
            
            last_update_time = current_time
            last_update_bytes = bytes_downloaded
            
            # 调用回调函数
            if progress_callback:
                try:
                    progress_callback(bytes_downloaded, total_size, speed)
                except InterruptedError:
                    # GUI请求停止下载
                    raise
        
        # 更新命令行进度条
        if use_cli_progress:
            pbar.update(chunk_size_downloaded)
    
    # 创建命令行进度条（如果没有进度回调）
    pbar = None
    if use_cli_progress:
        pbar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True, desc="Download")
    
    progress_lock = Lock()
    
    # Create a file with the correct size
    with open(filename, "wb") as f:
        f.truncate(total_size)
    
    if supports_ranges and threads > 1:
        if use_cli_progress:
            click.secho(f"Server supports range requests, using {threads} threads for download", fg="green")
        
        # Calculate chunk size for each thread
        chunk_size_per_thread = total_size // threads
        
        def download_chunk(start, end, thread_id):
            headers = {"Range": f"bytes={start}-{end}"}
            
            try:
                with session.get(url, headers=headers, stream=True, timeout=timeout) as response:
                    response.raise_for_status()
                    
                    # Write the chunk to the correct position in the file
                    with open(filename, "r+b") as f:
                        f.seek(start)
                        for chunk in response.iter_content(chunk_size=chunk_size):
                            if chunk:
                                f.write(chunk)
                                with progress_lock:
                                    update_progress(len(chunk))
            except Exception as e:
                if use_cli_progress:
                    click.secho(f"Thread {thread_id} failed: {e}", fg="red")
                return False
            except InterruptedError:
                # 传递停止信号
                raise
            
            return True
        
        # Create and start threads
        futures = []
        success = True
        
        try:
            with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
                for i in range(threads):
                    start = i * chunk_size_per_thread
                    # The last thread should download to the end of the file
                    end = total_size - 1 if i == threads - 1 else (i + 1) * chunk_size_per_thread - 1
                    futures.append(executor.submit(download_chunk, start, end, i))
                
                # Check if all threads completed successfully
                for future in concurrent.futures.as_completed(futures):
                    try:
                        result = future.result()
                        if not result:
                            success = False
                    except InterruptedError:
                        # 取消所有进行中的任务
                        for future in futures:
                            future.cancel()
                        raise
                    except Exception as e:
                        if use_cli_progress:
                            click.secho(f"Thread failed: {e}", fg="red")
                        success = False
                
                if not success:
                    raise Exception("Some download threads failed")
        except Exception as e:
            # 确保清理进度条
            if pbar:
                pbar.close()
            raise
                
    else:
        if use_cli_progress:
            if not supports_ranges:
                click.secho("Server does not support range requests, falling back to single thread download", fg="yellow")
            elif threads <= 1:
                click.secho("Single thread download mode", fg="yellow")
        
        # Fallback to single thread download
        try:
            with session.get(url, stream=True, timeout=timeout) as response:
                response.raise_for_status()
                with open(filename, "wb") as f:
                    for chunk in response.iter_content(chunk_size=chunk_size):
                        if chunk:
                            f.write(chunk)
                            update_progress(len(chunk))
        except InterruptedError:
            # 传递停止信号
            raise
    
    # 关闭进度条
    if pbar:
        pbar.close()
    
    # 最终进度更新
    if progress_callback:
        final_speed = (bytes_downloaded / (time.time() - start_time)) / (1024 * 1024)
        try:
            progress_callback(bytes_downloaded, total_size, final_speed)
        except InterruptedError:
            pass

def download_smb(url, filename, chunk_size, username, password, timeout, progress_callback=None):
    """Download file using SMB protocol with pysmb - supports multi-threaded chunked download"""
    # Parse SMB URL format: smb://[user:pass@]server/share/path/to/file
    parsed_url = urlparse(url)
    
    # 提取服务器名（可能包含端口号）
    netloc = parsed_url.netloc
    
    # 如果netloc包含@，则提取用户信息和服务器名
    if '@' in netloc:
        auth_part, server_part = netloc.split('@', 1)
        # 如果未提供用户名但URL中包含，则使用URL中的用户名
        if not username and ':' in auth_part:
            url_username, url_password = auth_part.split(':', 1)
            if not username:
                username = url_username
            if not password:
                password = url_password
        elif not username:
            username = auth_part
        server = server_part
    else:
        server = netloc
    
    # 分离服务器名和端口号
    server_host = server
    server_port = 445
    if ':' in server:
        server_host, port_str = server.split(':', 1)
        try:
            server_port = int(port_str)
        except ValueError:
            click.secho(f"Invalid port number: {port_str}, using default 445", fg="yellow")
    
    path_parts = parsed_url.path.strip('/').split('/')
    
    if len(path_parts) < 2:
        raise ValueError("SMB URL must be in format: smb://server/share/path/to/file")
    
    share = path_parts[0]
    file_path = '/'.join(path_parts[1:])
    
    # Use current user if no credentials provided
    if not username:
        username = os.environ.get('USERNAME', '')
    
    # 使用客户端机器名（pysmb需要这个参数）
    client_machine_name = os.environ.get('COMPUTERNAME', 'txpq-client')
    
    # 初始化变量
    conn = None
    total_size = 0
    bytes_read = 0
    
    # 是否需要显示命令行进度条
    use_cli_progress = progress_callback is None
    pbar = None
    
    # 初始化进度跟踪变量
    bytes_downloaded = 0
    start_time = time.time()
    last_update_time = start_time
    last_update_bytes = 0
    
    # 进度更新函数
    def update_progress(chunk_size_downloaded):
        nonlocal bytes_downloaded, last_update_time, last_update_bytes
        bytes_downloaded += chunk_size_downloaded
        
        # 计算下载速度
        current_time = time.time()
        time_diff = current_time - last_update_time
        
        # 每0.5秒或下载了足够的数据时更新
        if time_diff > 0.5 or bytes_downloaded >= total_size:
            bytes_diff = bytes_downloaded - last_update_bytes
            speed = (bytes_diff / time_diff) / (1024 * 1024)  # MB/s
            
            last_update_time = current_time
            last_update_bytes = bytes_downloaded
            
            # 调用回调函数
            if progress_callback:
                try:
                    progress_callback(bytes_downloaded, total_size, speed)
                except InterruptedError:
                    # GUI请求停止下载
                    raise
        
        # 更新命令行进度条
        if use_cli_progress and pbar:
            pbar.update(chunk_size_downloaded)
    
    try:
        # 连接到SMB服务器
        if use_cli_progress:
            click.secho(f"Connecting to SMB server: {server_host}:{server_port}", fg="blue")
        conn = SMBConnection(username, password, client_machine_name, server_host, use_ntlm_v2=True)
        conn.connect(server_host, server_port, timeout=timeout)
        if use_cli_progress:
            click.secho(f"Successfully connected to server: {server_host}", fg="green")
        
        # 获取文件信息和大小
        if use_cli_progress:
            click.secho(f"Getting file information from share: {share}, path: {file_path}", fg="blue")
        file_attributes, file_size = get_smb_file_info(conn, share, file_path)
        total_size = file_size
        if use_cli_progress:
            click.secho(f"File size: {total_size:,} bytes", fg="green")
        
        # 创建进度条
        if use_cli_progress:
            pbar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True, desc="Download")
        
        progress_lock = Lock()
        
        # 记录开始时间
        start_time = time.time()
        
        # 创建目标文件并设置大小
        with open(filename, "wb") as f:
            f.truncate(total_size)
        
        # 实现多线程分块下载
        threads = 4  # 使用4个线程，可以根据需要调整
        if use_cli_progress:
            click.secho(f"Starting multi-threaded download with {threads} threads", fg="green")
        
        # 计算每个线程的块大小
        chunk_size_per_thread = total_size // threads
        
        # 定义单个块下载函数
        def download_chunk(start, end, thread_id):
            thread_conn = None
            try:
                # 每个线程创建自己的连接
                thread_conn = SMBConnection(username, password, f"{client_machine_name}-thread{thread_id}", 
                                          server_host, use_ntlm_v2=True)
                thread_conn.connect(server_host, server_port, timeout=timeout)
                
                # 下载指定范围的数据
                bytes_downloaded = 0
                current_pos = start
                
                # 打开本地文件进行写入
                with open(filename, "r+b") as f:
                    f.seek(start)
                    
                    while current_pos < end:
                        # 计算本次读取的块大小
                        read_size = min(chunk_size, end - current_pos + 1)
                        
                        # 读取SMB文件数据
                        thread_conn.retrieveFileFromOffset(share, file_path, f, current_pos, read_size)
                        
                        with progress_lock:
                            update_progress(read_size)
                        
                        bytes_downloaded += read_size
                        current_pos += read_size
                
                thread_conn.close()
                if use_cli_progress:
                    click.secho(f"Thread {thread_id} completed download of {bytes_downloaded:,} bytes", fg="green")
                return True
            except InterruptedError:
                # 传递停止信号
                if thread_conn:
                    thread_conn.close()
                raise
            except Exception as e:
                if use_cli_progress:
                    click.secho(f"Thread {thread_id} failed: {e}", fg="red")
                if thread_conn:
                    thread_conn.close()
                return False
        
        # 创建并启动线程
        futures = []
        success = True
        
        try:
            with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
                for i in range(threads):
                    start = i * chunk_size_per_thread
                    # 最后一个线程下载到文件末尾
                    end = total_size - 1 if i == threads - 1 else (i + 1) * chunk_size_per_thread - 1
                    futures.append(executor.submit(download_chunk, start, end, i))
                
                # 检查所有线程是否成功完成
                for future in concurrent.futures.as_completed(futures):
                    try:
                        result = future.result()
                        if not result:
                            success = False
                    except InterruptedError:
                        # 取消所有进行中的任务
                        for future in futures:
                            future.cancel()
                        raise
                    except Exception as e:
                        if use_cli_progress:
                            click.secho(f"Thread failed: {e}", fg="red")
                        success = False
                
                if not success:
                    raise Exception("Some download threads failed")
        except Exception as e:
            # 确保清理进度条
            if pbar:
                pbar.close()
            raise
        
        # 关闭进度条
        if pbar:
            pbar.close()
        
        # 最终进度更新
        if progress_callback:
            final_speed = (bytes_downloaded / (time.time() - start_time)) / (1024 * 1024)
            try:
                progress_callback(bytes_downloaded, total_size, final_speed)
            except InterruptedError:
                pass
        
        # 计算下载时间
        time_cost = time.time() - start_time
        if use_cli_progress:
            click.secho(f"Download completed, time cost: {time_cost:.2f} seconds", fg="green")
            click.secho(f"Total bytes written: {total_size}", fg="blue")
            
            # 计算传输速度
            if time_cost > 0:
                speed_mb = (total_size / (1024 * 1024)) / time_cost
                click.secho(f"Average download speed: {speed_mb:.2f} MB/s", fg="green")
                
        # 验证文件大小
        local_file_size = os.path.getsize(filename)
        if local_file_size == total_size:
            if use_cli_progress:
                click.secho("File size verification successful", fg="green")
        else:
            if use_cli_progress:
                click.secho(f"Warning: File size mismatch! Expected {total_size}, got {local_file_size}", fg="red")
            raise ValueError(f"File size mismatch! Expected {total_size}, got {local_file_size}")
            
    except Exception as e:
        if use_cli_progress:
            click.secho(f"Error during SMB download: {e}", fg="red")
        raise
    finally:
        # 清理连接
        if use_cli_progress:
            click.secho("Cleaning up connections...", fg="blue")
        try:
            if conn is not None:
                conn.close()
        except:
            pass
        if use_cli_progress:
            click.secho("Connection cleanup completed", fg="green")

def get_smb_file_info(connection, share, file_path):
    """获取SMB文件的信息和大小"""
    try:
        # 列出文件信息
        file_info = connection.getAttributes(share, file_path)
        return file_info, file_info.file_size
    except Exception as e:
        # 如果getAttributes失败，尝试使用listPath获取信息
        try:
            # 获取文件所在目录
            dir_path = '/'.join(file_path.split('/')[:-1]) or '/'  # 处理根目录情况
            filename = file_path.split('/')[-1]
            
            # 列出目录中的文件
            files = connection.listPath(share, dir_path)
            for f in files:
                if f.filename == filename:
                    return f, f.file_size
            
            raise FileNotFoundError(f"File not found: {file_path}")
        except Exception as list_error:
            click.secho(f"Failed to get file info with getAttributes, also failed with listPath: {list_error}", fg="red")
            raise Exception(f"Unable to get file information: {str(e)}")

if __name__ == "__main__":
    download()